package com.example.rtsp.handler;
import com.example.rtsp.service.FFmpegService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.BinaryWebSocketHandler;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.io.IOException;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @Author: 杜科璇
 * @Description: 视频流处理程序
 * @Date: 2025/9/5 17:57
 * @Version: 1.0
 */


@Component
public class VideoStreamHandler extends BinaryWebSocketHandler {
    private final FFmpegService ffmpegService; // 注入FFmpegService
    private static final Logger logger = LoggerFactory.getLogger(VideoStreamHandler.class); // 日志记录

    private final Set<WebSocketSession> sessions = new CopyOnWriteArraySet<>(); // 用于存储WebSocket会话
    // 控制是否应该传输视频帧的标志
    private final AtomicBoolean shouldStream = new AtomicBoolean(false);

    // 控制是否暂停的标志
    private final AtomicBoolean isPaused = new AtomicBoolean(false);


    // 改为构造器注入
    @Autowired
    public VideoStreamHandler(FFmpegService ffmpegService) {
        this.ffmpegService = ffmpegService;
    }

    @PostConstruct
    public void init() {
        ffmpegService.setFrameConsumer(this::broadcastFrame); // 设置FFmpeg服务的帧消费者
//        ffmpegService.start();
    }

    @PreDestroy
    public void cleanup() {
        ffmpegService.stop(); // 停止FFmpeg服务
    }
    // 添加开始播放视频流的方法
    public void startStreaming() {
        if (shouldStream.compareAndSet(false, true)) {
            logger.info("开始视频流传输");
            ffmpegService.start();
        }
    }

    // 添加停止播放视频流的方法
    public void stopStreaming() {
        if (shouldStream.compareAndSet(true, false)) {
            logger.info("停止视频流传输");
            ffmpegService.stop();
        }
    }

    // 添加暂停视频流的方法
    public void pauseStreaming() {
        if (shouldStream.get() && !isPaused.get()) {
            isPaused.set(true);
            logger.info("暂停视频流传输");
        }
    }

    // 添加恢复视频流的方法
    public void resumeStreaming() {
        if (shouldStream.get() && isPaused.get()) {
            isPaused.set(false);
            logger.info("恢复视频流传输");
        }
    }


    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        // 只有在应该传输视频流时才接受新连接
        if (shouldStream.get()) {
            sessions.add(session);
            logger.info("新客户端连接: {}", session.getId());
        } else {
            // 如果不应该传输视频流，立即关闭连接
            logger.info("拒绝客户端连接，视频流未启动: {}", session.getId());
            try {
                session.close(CloseStatus.NOT_ACCEPTABLE);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }


    // 处理WebSocket连接关闭事件
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        logger.info("WebSocket连接已关闭: {}, 状态: {}", session.getId(), status);
        sessions.remove(session);

        // 检查是否还有其他活跃连接
        if (sessions.isEmpty()) {
            logger.info("没有活跃的客户端连接");
            // 只有当我们正在传输视频流时才停止FFmpeg
            if (shouldStream.get()) {
                logger.info("停止FFmpeg服务");
                stopStreaming();
            }
        }

        super.afterConnectionClosed(session, status); // 调用父类的方法
    }

    // 处理视频帧数据
    private void broadcastFrame(byte[] frame) {
        // 只有在shouldStream为true时才发送帧数据
        if (!shouldStream.get() || sessions.isEmpty()) return;

//        sessions.forEach(session -> {
        sessions.parallelStream().forEach(session -> {
            try {
                if (session.isOpen()) {
                    // 为每个会话创建独立的BinaryMessage实例，避免复用导致的问题
                    BinaryMessage message = new BinaryMessage(frame);
                    synchronized (session) {
                        session.sendMessage(message);
                    }
                } else {
                    sessions.remove(session);
                }
            } catch (IOException e) {
                logger.warn("发送帧到客户端失败: {}", e.getMessage());
                sessions.remove(session);
            }
        });
    }

    // 获取当前活跃会话数量
    public int getActiveSessionCount() {
        return sessions.size();
    }

    // 检查是否正在传输视频流
    public boolean isStreaming() {
        return shouldStream.get();
    }

    // 检查是否暂停
    public boolean isPaused() {
        return isPaused.get();
    }
}